from sqlalchemy.ext.compiler import compiles
import pandas as pd
import time
from sqlalchemy import create_engine, Table, MetaData
from sqlalchemy.sql.expression import Insert
from sqlalchemy.dialects.mysql import insert
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine, text
import platform
import json
import requests
import uuid
from urllib.parse import quote_plus
import traceback
from datetime import datetime
import re

def classify_exhibition_industry(exhibition_data):
    """
    调用AI接口仅对展会行业进行分类，使用流式输出
    """
    
    # 生成唯一的dataId和chatId
    dataId = str(uuid.uuid4()).replace("-", "")
    chatId = str(uuid.uuid4()).replace("-", "")
    
    print(exhibition_data)
    promot = f"""
    以下是展会信息，请你根据知识库查询，仅对展会行业类别进行分类并返回分类结果：
    {exhibition_data}
    
    请你一定返回json格式数据（其它数据不要)：
    {{
        'status': true/false,  # 分类是否成功
        'industry': ""  # 你对展会的行业分类结果
    }}
    
    注意：
    对于行业字段，请必须在以下行业中，选择一个最合适的来作为字段值：
        建筑建材
        综合类
        工业机械
        文体娱乐
        采矿及设备
        食品饮料
        纺织服饰
        冷暖设备
        农林牧渔
        电力及新能源设备
        船舶和轨道等交通
        珠宝首饰
        食品加工设备
        家具家居
        公共安全
        工程机械
        高新材料 / 技术
        节能环保
        医疗及制药设备
        钟表眼镜
        酒店用品
        其他设备 / 材料
        汽车
        礼品玩具及宠物用品
        仪表仪器
        皮革鞋包
        农业机械
        计算机通信及电子设备
        物流仓储
        自行车电动车摩托车
        其他
        汽配及服务
        印刷包装
        光电产品
        家电
        办公用品
        化工产品
        五金工具
        美容美发
        广告传播
        工程机械及管材线材
        妇婴用品
        泵阀管道
        电商和交易会
        包装机械
        航空航天
        电梯
        药品及原料
        电力电工
        光电
        纺织及制衣机械
        食品及饮料
        物流与供应链
        医疗实验室设备
        建筑、建材及施工
        涂料、油漆及表面处理
    请仅返回JSON格式数据！！，不要附加任何额外解释！
    """

    # 构建提示信息
    messages = [{
        "dataId": dataId,
        "hideInUI": False,
        "role": "user",
        "content": f"{promot}"
    }]
    
    # 使用新的API密钥和Bearer认证
    API_KEY = "fastgpt-kspcWF3r4XnI9leHXcZqE6tkY2LKGx7p9WgT3fMD9Ywtn06zyQ36oFdo4Vb"
    
    # 使用新的API接口和请求格式
    url = "http://106.53.97.117:3000/api/v1/chat/completions"
    
    payload = {
        "chatId": chatId,
        "appId": "67e014c77a46e4ccc286d102",
        "variables": {
            "cTime": time.strftime("%Y-%m-%d %H:%M:%S %A")
        },
        "messages": [
            {
                "content": promot,
                "role": "user"
            }
        ],
        "modules": [
            {
                "moduleId": "dataset",
                "datasetId": "67e3b1d27a46e4ccc2879bfd"  # 使用正确的数据集ID
            }
        ],
        "stream": True
    }
    
    headers = {
        "Accept": "text/event-stream",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}",  # 使用Bearer令牌认证
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36"
    }
    
    # 添加重试机制
    max_retries = 3
    retry_delay = 2  # 秒
    last_exception = None
    
    for retry_count in range(max_retries):
        try:
            # 增加超时设置
            response = requests.post(
                url, 
                json=payload, 
                headers=headers, 
                timeout=(10, 300),  # 连接超时10秒，读取超时300秒
                stream=True
            )
            response.raise_for_status()
            
            full_response = ""
            print(f"API响应状态码: {response.status_code}")
            
            for line in response.iter_lines():
                if line:
                    try:
                        # 移除 "data: " 前缀并解析 JSON
                        if line.startswith(b'data: '):
                            json_str = line[6:].decode('utf-8')
                            
                            # 处理结束标记
                            if json_str.strip() == '[DONE]':
                                print("收到完成标记")
                                break
                                
                            # 解析JSON数据
                            try:
                                data = json.loads(json_str)
                                
                                # 提取内容 - 新的响应格式
                                if 'data' in data and 'text' in data['data']:
                                    content = data['data']['text']
                                    print(content, end='', flush=True)
                                    full_response += content
                                # 兼容旧格式
                                elif 'choices' in data and len(data['choices']) > 0:
                                    content = data['choices'][0].get('delta', {}).get('content', '')
                                    if content:
                                        print(content, end='', flush=True)
                                        full_response += content
                            except json.JSONDecodeError:
                                # 如果无法解析为JSON，跳过
                                continue
                    except Exception as e:
                        print(f"处理流式响应时出错: {str(e)}")
                        continue
            
            print(f"\n完整响应: {full_response[:200]}...")
            
            # 验证响应是否为空
            if not full_response.strip():
                print("警告: 收集到的完整响应为空")
                if retry_count < max_retries - 1:
                    wait_time = retry_delay * (2 ** retry_count)
                    print(f"等待 {wait_time} 秒后重试...")
                    time.sleep(wait_time)
                    continue
                return {"status": False, "error": "API返回了空响应内容"}
                
            return {"status": True, "response": full_response}
            
        except requests.exceptions.RequestException as e:
            last_exception = e
            print(f"API请求失败 (尝试 {retry_count+1}/{max_retries}): {str(e)}")
            
            # 如果不是最后一次尝试，则等待后重试
            if retry_count < max_retries - 1:
                wait_time = retry_delay * (2 ** retry_count)  # 指数退避
                print(f"等待 {wait_time} 秒后重试...")
                time.sleep(wait_time)
            else:
                print(f"达到最大重试次数 ({max_retries})，放弃请求")
                return {"status": False, "error": str(e)}
    
    # 如果所有重试都失败
    return {"status": False, "error": str(last_exception) if last_exception else "所有重试都失败"}

class MySQLHelper:
    def __init__(self):
        if platform.system() == 'Linux':
            self.host = "mysql"
            self.port = 3306
            self.user = "root"
            self.password = quote_plus("123456@exhibition")
            self.db = "global_exhibition_matrix"
        else:
            self.host = "106.53.97.117"
            self.port = 3306
            self.user = "my_user"
            self.password = "my_password"
            self.db = "my_database"
        self.charset = "utf8mb4"
        self.engine = None
        self.Session = None
        self._create_engine()

    def _create_engine(self):
        # 创建数据库连接引擎，使用连接池
        connection_string = f"mysql+pymysql://{self.user}:{self.password}@{self.host}:{self.port}/{self.db}?charset={self.charset}"
        self.engine = create_engine(connection_string, pool_size=10, max_overflow=1490, pool_recycle=3600)
        self.Session = sessionmaker(bind=self.engine)

    def get_session(self):
        # 获取数据库会话
        return self.Session()

    def execute_select(self, sql):
        time_start = time.time()
        with self.engine.connect() as conn:
            # 使用 text() 函数将 SQL 语句转换为可执行对象
            sql_text = text(sql)
            result = conn.execute(sql_text)
            data = []
            for row in result:
                # 使用 row._mapping 将 RowProxy 转换为字典
                data.append(dict(row._mapping))
            time_end = time.time()
            # print(f"执行时间：{time_end - time_start}")
            return data

    def execute_edit(self, sql):
        with self.engine.connect() as conn:
            result = conn.execute(text(sql))
            conn.commit()
            return result.rowcount


db = MySQLHelper()

fields_mapping = {
    'chinese_name': '中文名称',
    'english_name': '英文全称',
    'english_short_name': '英文简称',
    'logo': '标识',
    'holding_time': '举办时间',
    'event_date_start': '举办开始日期',
    'event_date_end': '举办结束日期',
    'event_month': '举办月份',
    'continent': '洲',
    'country': '国别',
    'city': '城市',
    'frequency': '举办频率',
    'industry': '行业',
    'organizer': '主办方',
    'organizer_en': '主办方英文',
    'contact_info': '联系方式',
    'exhibition_hall': '展馆',
    'exhibition_hall_english': '展馆英文',
    'official_website': '官网',
    'exhibit_scope': '展品范围',
    'description': '简介'
}

def get_current_industry_status():
    """
    获取当前展会行业分类状态统计
    """
    try:
        # 查询各状态的展会数量
        status_sql = """
            SELECT 
                SUM(CASE WHEN industry IS NULL OR industry = '' THEN 1 ELSE 0 END) as missing_industry_count,
                SUM(CASE WHEN industry IS NOT NULL AND industry != '' THEN 1 ELSE 0 END) as classified_count,
                COUNT(*) as total_count
            FROM exhibitions
        """
        print(status_sql,"sql语句")
        
        result = db.execute_select(status_sql)
        if result and len(result) > 0:
            status_data = result[0]
            
            print("\n" + "="*50)
            print("当前展会行业分类状态统计")
            print("="*50)
            print(f"未分类展会: {status_data.get('missing_industry_count', 0)} 条")
            print(f"已分类展会: {status_data.get('classified_count', 0)} 条")
            print(f"总展会记录数: {status_data.get('total_count', 0)} 条")
            print("="*50)
            
            return status_data
        else:
            print("获取当前状态统计失败")
            return None
    except Exception as e:
        print(f"获取当前状态统计时出错: {str(e)}")
        print(traceback.format_exc())
        return None

def main():
    try:
        print("="*50)
        print("开始处理展会行业分类")
        print("="*50)
        
        # 获取处理前的状态统计
        print("处理前状态统计:")
        get_current_industry_status()
        
        # 查询需要处理的展会记录
        # 选择未分类或行业字段为空的展会
        fields = 'id, chinese_name, english_name, exhibit_scope, description'
        
        sql = f"""
            SELECT {fields}
            FROM exhibitions 
            WHERE (industry_ai IS NULL OR industry_ai = '')
        """
        
        # 尝试多次查询数据库，防止临时连接问题
        max_db_retries = 13
        result_json = None
        
        for retry in range(max_db_retries):
            try:
                result_json = db.execute_select(sql)
                break
            except Exception as e:
                print(f"数据库查询失败 (尝试 {retry+1}/{max_db_retries}): {str(e)}")
                if retry < max_db_retries - 1:
                    print("等待 5 秒后重试...")
                    time.sleep(5)
                else:
                    print("达到最大重试次数，无法查询数据库")
                    raise
        
        if not result_json or len(result_json) == 0:
            print("没有需要处理的展会记录")
            return
            
        result_df = pd.DataFrame(result_json)

        # 处理所有记录
        total_records = len(result_df)
        print(f"总共需要处理 {total_records} 条展会记录")
        
        success_count = 0
        failed_count = 0
        start_time = time.time()

        for index, row in result_df.iterrows():
            try:
                exhibition_data = {
                    "展会ID": row['id'],
                    "中文名称": row['chinese_name'] if 'chinese_name' in row and row['chinese_name'] else "",
                    "英文全称": row['english_name'] if 'english_name' in row and row['english_name'] else "",
                    "展品范围": row['exhibit_scope'] if 'exhibit_scope' in row and row['exhibit_scope'] else "",
                    "简介": row['description'] if 'description' in row and row['description'] else ""
                }
                
                exhibition_id = row['id']
                
                print(f"\n[{index+1}/{total_records}] 正在处理第 {index+1} 条记录")
                print(f"展会ID: {exhibition_id}")
                print(f"展会名称: {exhibition_data.get('中文名称', '未知')}")
                
                # 调用AI接口进行行业分类
                try:
                    classify_result = classify_exhibition_industry(exhibition_data)
                    print(f"AI处理状态: {'成功' if classify_result['status'] else '失败'}")
                except Exception as e:
                    print(f"调用AI接口时出错: {str(e)}")
                    failed_count += 1
                    _update_classification_failed(exhibition_id, f"AI接口调用失败: {str(e)[:100]}")
                    continue
                
                if classify_result['status']:
                    try:
                        # 从response中提取数据
                        response_text = classify_result['response']
                        
                        # 尝试从响应中提取JSON
                        json_data = None
                        
                        # 尝试各种方法解析JSON数据
                        # 1. 如果响应包含markdown代码块，提取其中的JSON
                        if '```json' in response_text:
                            json_start = response_text.find('```json') + 7
                            if json_start == 6:  # 如果没有找到```json，尝试不带语言标记的版本
                                json_start = response_text.find('```') + 3
                            json_end = response_text.rfind('```')
                            if json_start != -1 and json_end != -1 and json_start < json_end:
                                json_str = response_text[json_start:json_end].strip()
                                try:
                                    json_data = json.loads(json_str)
                                except:
                                    print("从Markdown代码块提取JSON失败")
                        
                        # 2. 如果上面的方法失败，尝试在文本中搜索JSON对象
                        if not json_data:
                            json_matches = re.findall(r'({[\s\S]*?})', response_text)
                            for match in json_matches:
                                try:
                                    # 修复: 处理JavaScript风格的JSON字符串
                                    # 替换JavaScript布尔值
                                    fixed_match = match.replace('true', 'True').replace('false', 'False')
                                    # 使用eval安全地评估为Python字典
                                    # 添加安全检查，确保只有字典结构被评估
                                    if re.match(r'^{.*}$', fixed_match.strip()):
                                        potential_json = eval(fixed_match, {"__builtins__": {}}, {})
                                    else:
                                        raise ValueError("不安全的JSON格式")
                                    # 验证是否包含我们需要的字段
                                    if 'industry' in potential_json:
                                        json_data = potential_json
                                        break
                                except:
                                    continue
                        
                        # 3. 如果仍然没有找到有效JSON，尝试直接解析整个响应
                        if not json_data:
                            try:
                                cleaned_response = response_text.strip()
                                first_brace = cleaned_response.find('{')
                                last_brace = cleaned_response.rfind('}')
                                
                                if first_brace != -1 and last_brace != -1 and first_brace < last_brace:
                                    potential_json_str = cleaned_response[first_brace:last_brace+1]
                                    # 修复: 处理JavaScript风格的JSON字符串
                                    fixed_json_str = potential_json_str.replace('true', 'True').replace('false', 'False')
                                    try:
                                        # 添加安全检查
                                        if re.match(r'^{.*}$', fixed_json_str.strip()):
                                            json_data = eval(fixed_json_str, {"__builtins__": {}}, {})  # 使用eval代替json.loads
                                        else:
                                            raise ValueError("不安全的JSON格式")
                                    except:
                                        print("使用eval解析JSON失败")
                                else:
                                    # 修复: 处理JavaScript风格的JSON字符串
                                    fixed_response = cleaned_response.replace('true', 'True').replace('false', 'False')
                                    # 添加安全检查
                                    if re.match(r'^{.*}$', fixed_response.strip()):
                                        json_data = eval(fixed_response, {"__builtins__": {}}, {})  # 使用eval代替json.loads
                                    else:
                                        raise ValueError("不安全的JSON格式")
                            except:
                                print("直接解析整个响应失败")
                                raise json.JSONDecodeError("无法从响应中提取有效JSON", response_text, 0)
                        
                        # 确保JSON数据包含必要的字段
                        if not json_data:
                            print("无法从AI响应中解析出有效的JSON数据")
                            raise ValueError("无法从AI响应中解析出有效的JSON数据")
                        
                        # 提取行业信息
                        industry = json_data.get('industry', '')
                        
                        # 确保行业不为空
                        if not industry:
                            print("AI未返回有效的行业分类")
                            failed_count += 1
                            _update_classification_failed(exhibition_id, "AI未返回有效的行业分类")
                            continue
                        
                        # SQL转义
                        safe_industry = industry.replace("'", "''")
                        
                        # 更新数据库，仅更新industry字段
                        update_sql = f"""
                        UPDATE exhibitions 
                        SET industry_ai = '{safe_industry}',industry = '{safe_industry}'
                        WHERE id = {exhibition_id}
                        """
                        
                        # 执行更新并添加重试机制
                        update_success = False
                        max_update_retries = 3
                        
                        for update_retry in range(max_update_retries):
                            try:
                                db.execute_edit(update_sql)
                                update_success = True
                                break
                            except Exception as e:
                                print(f"数据库更新失败 (尝试 {update_retry+1}/{max_update_retries}): {str(e)}")
                                if update_retry < max_update_retries - 1:
                                    print("等待 3 秒后重试更新...")
                                    time.sleep(3)
                        
                        if update_success:
                            print(f"成功更新行业分类: {industry}")
                            success_count += 1
                        else:
                            print("达到最大重试次数，无法更新数据库")
                            failed_count += 1
                            continue
                        
                    except json.JSONDecodeError as e:
                        print(f"JSON解析错误: {str(e)}")
                        failed_count += 1
                        _update_classification_failed(exhibition_id, f"JSON解析错误: {str(e)[:100]}")
                        continue
                    except Exception as e:
                        print(f"处理AI分类结果时出错: {str(e)}")
                        print(traceback.format_exc())
                        failed_count += 1
                        _update_classification_failed(exhibition_id, f"处理错误: {str(e)[:100]}")
                        continue
                else:
                    print(f"AI分类失败: {classify_result.get('error', '未知错误')}")
                    failed_count += 1
                    _update_classification_failed(exhibition_id, f"AI分类失败: {classify_result.get('error', '未知错误')[:100]}")
                    continue
                
                # 每处理10条记录显示一次进度统计
                if (index + 1) % 10 == 0 or (index + 1) == total_records:
                    elapsed_time = time.time() - start_time
                    avg_time = elapsed_time / (index + 1)
                    remaining_count = total_records - (index + 1)
                    est_remaining_time = avg_time * remaining_count
                    
                    print(f"\n--- 处理进度 ---")
                    print(f"已处理: {index+1}/{total_records} ({(index+1)/total_records*100:.1f}%)")
                    print(f"成功: {success_count}, 失败: {failed_count}")
                    print(f"已用时间: {elapsed_time:.1f}秒")
                    if remaining_count > 0:
                        print(f"预计剩余时间: {est_remaining_time:.1f}秒 (约 {est_remaining_time/60:.1f}分钟)")
                
                # 添加短暂延时，避免请求过于频繁
                time.sleep(2)
                
            except Exception as e:
                print(f"处理记录时发生错误: {str(e)}")
                print(traceback.format_exc())
                failed_count += 1
                try:
                    _update_classification_failed(exhibition_id, f"处理错误: {str(e)[:100]}")
                except:
                    print("更新失败状态时也发生错误")
                continue
        
        # 最终统计
        total_time = time.time() - start_time
        print("\n"+"="*50)
        print("处理完成统计")
        print("="*50)
        print(f"总记录数: {total_records}")
        print(f"成功处理: {success_count} ({success_count/total_records*100:.1f}%)")
        print(f"处理失败: {failed_count} ({failed_count/total_records*100:.1f}%)")
        print(f"总耗时: {total_time:.1f}秒 (约 {total_time/60:.1f}分钟)")
        print("="*50)
        
        # 重新查询一次当前状态，显示实时数据
        print("\n处理后状态统计:")
        get_current_industry_status()
        
        # 检查是否还有待处理的记录
        current_status = get_current_industry_status()
        if current_status and current_status.get('missing_industry_count', 0) > 0:
            print(f"\n仍有 {current_status.get('missing_industry_count')} 条记录未分类")
            print("您可以再次运行此脚本继续处理剩余记录")
        
    except Exception as e:
        print(f"程序执行过程中发生错误: {str(e)}")
        print(traceback.format_exc())
    finally:
        print("\n所有展会行业分类处理完成！")

def _update_classification_failed(exhibition_id, reason):
    """记录行业分类失败的原因"""
    try:
        safe_reason = reason.replace("'", "''")  # SQL转义
        # 修复: 只使用industry_ai_error字段，不使用industry_update_time
        # 尝试使用industry_ai_error字段
        update_sql = f"""
        UPDATE exhibitions 
        SET industry_ai_error = '分类失败: {safe_reason}'
        WHERE id = {exhibition_id}
        """
        try:
            db.execute_edit(update_sql)
        except Exception as e:
            # 如果industry_ai_error字段不存在，尝试记录到日志而不是数据库
            if "Unknown column 'industry_ai_error'" in str(e):
                print(f"注意: 数据库表缺少'industry_ai_error'列，无法记录错误信息。错误信息: {reason}")
                # 替代方案：更新industry_ai字段
                update_sql_alt = f"""
                UPDATE exhibitions 
                SET industry_ai = '分类失败: {safe_reason}'
                WHERE id = {exhibition_id}
                """
                db.execute_edit(update_sql_alt)
            else:
                # 其他错误
                raise e
        
        print(f"展会ID {exhibition_id} 行业分类失败: {reason}")
        return True
    except Exception as e:
        print(f"更新展会ID {exhibition_id} 行业分类失败状态时出错: {str(e)}")
        return False

if __name__ == "__main__":
    main() 